from dataclasses import dataclass, field
from typing import List
import pandas as pd
from . import base


@dataclass
class Select(base.BaseTransformer, base.CopyMixin):
    cols: List[str] = field(default_factory=lambda: ['close'])
    
    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        return data[self.cols]
    
    
@dataclass
class Drop(base.BaseTransformer, base.CopyMixin):
    cols: List[str] = field(default_factory=lambda: ['close'])
    
    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        return data.drop(columns=self.cols)
